package test3;

/**
 * @Author Jun jie Yi
 * 2022/4/24 18:43
 */
public class Demo5 {

    // k 个桶（集合），记录每个桶装的数字之和
    public static void main(String[] args) {
        //将 nums 平均分到 k 个桶里面
        int[] nums = {1, 2, 3, 2, 3, 5, 4, 3, 2, 4, 5, 7};
        boolean b = canPartitionKSubsets(nums, 4);
        System.out.println(b);
    }

    // 穷举 nums 中的每个数字
    static void backtrack(int[] nums, int index, int[] bucket) {
        // base case
        if (index == nums.length) {
            return;
        }
        // 穷举每个桶
        for (int i = 0; i < bucket.length; i++) {
            // 选择装进第 i 个桶
            bucket[i] += nums[index];
            // 递归穷举下一个数字的选择
            backtrack(nums, index + 1, bucket);
            // 撤销选择
            bucket[i] -= nums[index];
        }
    }

    // 主函数
    static boolean canPartitionKSubsets(int[] nums, int k) {
        // 排除一些基本情况
        if (k > nums.length) return false;
        int sum = 0;
        for (int v : nums) sum += v;
        if (sum % k != 0) return false;
        // k 个桶（集合），记录每个桶装的数字之和
        int[] bucket = new int[k];
        // 理论上每个桶（集合）中数字的和
        int target = sum / k;
        // 穷举，看看 nums 是否能划分成 k 个和为 target 的子集
        return backtrack(nums, 0, bucket, target);
    }

    // 递归穷举 nums 中的每个数字
    static boolean backtrack(int[] nums, int index, int[] bucket, int target) {
        if (index == nums.length) {
            // 检查所有桶的数字之和是否都是 target
            for (int i = 0; i < bucket.length; i++) {
                if (bucket[i] != target) {
                    return false;
                }
            }
            // nums 成功平分成 k 个子集
            return true;
        }

        // 穷举 nums[index] 可能装入的桶
        for (int i = 0; i < bucket.length; i++) {
            // 剪枝，桶装装满了
            if (bucket[i] + nums[index] > target) {
                continue;
            }
            // 如果桶没有装满
            // 将 nums[index] 装入 bucket[i]
            bucket[i] += nums[index];
            // 递归穷举下一个数字的选择
            if (backtrack(nums, index + 1, bucket, target)) {
                return true;
            }
            // 撤销选择
            bucket[i] -= nums[index];
        }
        // nums[index] 装入哪个桶都不行
        return false;
    }
}
